# -*- coding: utf-8 -*-
from typing import List, Dict, Any, Optional, Callable, Tuple
import numpy as np
from collections import defaultdict
import math

try:
    # 可选：使用 SciPy 的匈牙利算法更稳
    from scipy.optimize import linear_sum_assignment
    SCIPY_AVAILABLE = True
except Exception:
    SCIPY_AVAILABLE = False


# ========== 轻量 Pose 追踪器（稳定 person_id） ==========

# 使用 MediaPipe Pose 的关键点索引（与你的 common.landmarks 索引保持一致）
# 选取稳定度较高的上半身骨架进行匹配
MATCH_IDX = [
    0,   # NOSE
    11,  # LEFT_SHOULDER
    12,  # RIGHT_SHOULDER
    23,  # LEFT_HIP
    24,  # RIGHT_HIP
]

def _euclidean(a: Tuple[float, float], b: Tuple[float, float]) -> float:
    return math.sqrt((a[0]-b[0])**2 + (a[1]-b[1])**2)

def _lm_xyv(lm) -> Tuple[float, float, float]:
    # 将单个 landmark 转为 (x, y, visibility)，visibility 兜底为 1.0
    x = float(getattr(lm, "x", 0.0) or 0.0)
    y = float(getattr(lm, "y", 0.0) or 0.0)
    v = float(getattr(lm, "visibility", 1.0) if getattr(lm, "visibility", None) is not None else 1.0)
    return x, y, v

class _Track:
    def __init__(self, track_id: int, lm_list: List[Any], ema_alpha: float = 0.7):
        self.track_id = track_id
        self.lms = lm_list  # 原始 landmarks 列表（每个为 NormalizedLandmark）
        self.missed = 0
        self.ema_alpha = ema_alpha

    def update(self, new_lms: List[Any]):
        # 对每个 landmark 做 EMA 平滑
        a = self.ema_alpha
        if not self.lms or len(self.lms) != len(new_lms):
            # 长度不一致时，直接替换
            self.lms = new_lms
            self.missed = 0
            return
        merged = []
        for old_lm, new_lm in zip(self.lms, new_lms):
            ox, oy, ov = _lm_xyv(old_lm)
            nx, ny, nv = _lm_xyv(new_lm)
            try:
                old_lm.x = a * nx + (1 - a) * ox
                old_lm.y = a * ny + (1 - a) * oy
                oz = float(getattr(old_lm, "z", 0.0) or 0.0)
                nz = float(getattr(new_lm, "z", 0.0) or 0.0)
                old_lm.z = a * nz + (1 - a) * oz
                old_lm.visibility = a * nv + (1 - a) * ov
                merged.append(old_lm)
            except Exception:
                merged.append(new_lm)
        self.lms = merged
        self.missed = 0

class _PoseTracker:
    def __init__(self, max_cost: float = 0.15, max_missed: int = 10, ema_alpha: float = 0.7):
        self.max_cost = max_cost
        self.max_missed = max_missed
        self.ema_alpha = ema_alpha
        self._next_id = 0
        self._tracks: Dict[int, _Track] = {}

    def _distance(self, det_lms: List[Any], trk: _Track) -> float:
        pts = []
        for idx in MATCH_IDX:
            if idx >= len(det_lms) or idx >= len(trk.lms):
                continue
            x1, y1, v1 = _lm_xyv(det_lms[idx])
            x2, y2, v2 = _lm_xyv(trk.lms[idx])
            w = 0.5 * (max(0.05, v1) + max(0.05, v2))
            pts.append(w * _euclidean((x1, y1), (x2, y2)))
        if not pts:
            return 1e9
        return sum(pts) / len(pts)

    def _associate(self, detections: List[List[Any]]):
        track_ids = list(self._tracks.keys())
        if not track_ids or not detections:
            return [], list(range(len(detections))), track_ids

        cost_matrix = [
            [self._distance(detections[j], self._tracks[tid]) for j in range(len(detections))]
            for tid in track_ids
        ]

        if SCIPY_AVAILABLE:
            row_ind, col_ind = linear_sum_assignment(cost_matrix)
        else:
            # 贪心匹配作为回退
            triples = []
            for r, row in enumerate(cost_matrix):
                for c, cost in enumerate(row):
                    triples.append((cost, r, c))
            triples.sort(key=lambda x: x[0])
            used_r, used_c = set(), set()
            pairs = []
            for cost, r, c in triples:
                if r in used_r or c in used_c:
                    continue
                used_r.add(r); used_c.add(c)
                pairs.append((r, c))
            row_ind = [p[0] for p in pairs]
            col_ind = [p[1] for p in pairs]

        matches = []
        unmatched_det = set(range(len(detections)))
        unmatched_trk_rows = set(range(len(track_ids)))
        for r, c in zip(row_ind, col_ind):
            if cost_matrix[r][c] <= self.max_cost:
                matches.append((track_ids[r], c))
                unmatched_det.discard(c)
                unmatched_trk_rows.discard(r)

        unmatched_trk_ids = [track_ids[r] for r in unmatched_trk_rows]
        return matches, list(unmatched_det), unmatched_trk_ids

    def update(self, detections: List[List[Any]]) -> Dict[int, List[Any]]:
        matches, unmatched_det_idx, unmatched_trk_ids = self._associate(detections)

        # 更新匹配到的轨迹
        for tid, det_idx in matches:
            trk = self._tracks[tid]
            trk.ema_alpha = self.ema_alpha
            trk.update(detections[det_idx])

        # 新建轨迹
        for det_idx in unmatched_det_idx:
            tid = self._next_id
            self._next_id += 1
            self._tracks[tid] = _Track(track_id=tid, lm_list=detections[det_idx], ema_alpha=self.ema_alpha)

        # 未匹配轨迹 missed++
        to_del = []
        for tid in unmatched_trk_ids:
            trk = self._tracks.get(tid)
            if trk is None:
                continue
            trk.missed += 1
            if trk.missed > self.max_missed:
                to_del.append(tid)
        for tid in to_del:
            self._tracks.pop(tid, None)

        # 导出：当前活跃 person_id -> landmarks
        return {tid: trk.lms for tid, trk in self._tracks.items()}


# ========== 动作与管理类 ==========

class Action:
    label: str = "base"
    def score(self, landmarks, person_id: int, get_state: Callable[[int], dict]) -> Tuple[float, Dict[str, Any]]:
        """返回 0~1 的得分，允许使用 per-person 历史状态"""
        raise NotImplementedError

class ActionRecognitionManager:
    """
    - 维护动作列表
    - latest_poses: 每帧的人体关键点（稳定的 person_id -> landmarks）
    - states: per-person 的动作历史状态容器（dict），由各动作模块自行维护键空间
    - 集成追踪器，为每个人分配稳定唯一 person_id
    """
    def __init__(self):
        self.actions: List[Action] = []
        self.latest_poses: Dict[int, Any] = {}      # person_id -> list[landmark]（稳定 ID）
        self._states: Dict[int, dict] = defaultdict(dict)  # person_id -> state dict

        # 追踪器参数可按需调整
        self._tracker = _PoseTracker(max_cost=0.25, max_missed=10, ema_alpha=0.7)

    def register_action(self, action: Action):
        self.actions.append(action)

    def update_latest_poses(self, persons: List[Any]):
        """
        persons: List[List[landmark]] 来自当前帧的多人人体 landmarks 列表
        集成追踪器，输出稳定 ID 的 latest_poses，并维护 states 生命周期
        """
        tracked_map: Dict[int, List[Any]] = self._tracker.update(persons)
        self.latest_poses = tracked_map

        # 清理状态（对不再活跃的 person_id 进行清理）
        active_ids = set(tracked_map.keys())
        for pid in list(self._states.keys()):
            if pid not in active_ids:
                del self._states[pid]

    def get_latest_pose(self, person_id: int) -> Optional[Any]:
        return self.latest_poses.get(person_id)

    def get_all_latest_poses(self) -> Dict[int, Any]:
        """返回当前所有活跃人员的最新 landmarks"""
        return dict(self.latest_poses)

    def get_state(self, person_id: int) -> dict:
        """获取或创建 per-person 状态字典"""
        return self._states[person_id]

    def recognize_actions(self, landmarks, person_id: int) -> List[Dict[str, float]]:
        out = []
        state_getter = self.get_state
        for act in self.actions:
            try:
                result = act.score(landmarks, person_id, state_getter)
                if isinstance(result, tuple):
                    s, extra = result
                else:
                    s, extra = result, {}
                s = float(np.clip(s, 0.0, 1.0))
            except Exception:
                s, extra = 0.0, {}
            if s > 0:
                label = extra.get("label", getattr(act, "label", "swipe"))
                out.append({"label": label, "score": s})
        return out

    @staticmethod
    def select_top_actions(actions: List[Dict[str, float]], threshold: float, topk: int):
        filtered = [a for a in actions if a.get("score", 0.0) >= threshold]
        filtered.sort(key=lambda x: x["score"], reverse=True)
        return filtered[:topk]


class ActionRecognitionManagerRef:
    def __init__(self):
        self._ref: Optional[ActionRecognitionManager] = None

    def set(self, manager: ActionRecognitionManager):
        self._ref = manager

    def get(self) -> Optional[ActionRecognitionManager]:
        return self._ref